import os
import traceback
from typing import List

import numpy
import pandas
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter

from conf.config import BASE_DIR, logger


class ExcelSplitter(object):
    """
    Excel文本分割器
    excel表中要有“content”和“source”列
    """

    def __init__(self, excel_filepath, *args, **kwargs):
        super(ExcelSplitter, self).__init__(*args, **kwargs)
        self.excel_filepath = excel_filepath  # excel文件名
        self.docs: List[Document] = list()  # 保存文档
        self.chunks: List[Document] = list()  # 保存切割后的文本块

    def get_docs(self) -> List[Document]:
        """
        从excel文件中获取携带元数据的文档列表
        :return:
        """
        if not self.excel_filepath.endswith(".xlsx"):
            raise ValueError(f"excel文件必须是.xlsx后缀")

        logger.info(f"开始读取：{self.excel_filepath}")
        try:
            data_frame = pandas.read_excel(self.excel_filepath, index_col=None, usecols=['content', 'source'])
        except ValueError as e:
            error_str = traceback.format_exc()
            logger.error(error_str)
            raise ValueError(f"excel表中要有“content”和“source”列")

        doc_list = list(data_frame.values)
        for doc in doc_list:
            self.docs.append(Document(page_content=doc[0], metadata={'source': doc[1]}))

        return self.docs

    def get_chunks(self) -> List[Document]:
        """
        切割文本，得到文本块
        :param docs:
        :return:
        """
        text_splitter = RecursiveCharacterTextSplitter(
            separators=["。", "！", "？", "……", "!", "…", "\n"],  # "?" re.error: nothing to repeat at position 0
            chunk_size=250,
            chunk_overlap=0,
            length_function=len,
            keep_separator=False,
            add_start_index=True,
        )

        self.chunks = text_splitter.create_documents(
            texts=[doc.page_content for doc in self.docs],
            metadatas=[doc.metadata for doc in self.docs]
        )

        return self.chunks

    def save_chunks_to_csv(self):
        """
        保存文本块到csv文件
        :return:
        """
        csv_output_path = os.path.splitext(self.excel_filepath)[0] + ".csv"

        rows = [[chunk.page_content, chunk.metadata.get("source")] for chunk in self.chunks]

        header = ["text_chunk", "source"]
        df = pandas.DataFrame(rows, columns=header)  # 组成一个csv

        logger.info(f"开始写入：{csv_output_path}")
        df.index = numpy.arange(1, len(df)+1)  # 写入csv时，id从1开始
        df.to_csv(csv_output_path, index=True, index_label="text_chunk_id")


def main():
    """ 切割excel中的文本，并存为csv文件 """

    excel_filepath = os.path.join(BASE_DIR, f'docs/template.xlsx')
    excel_splitter = ExcelSplitter(
        excel_filepath=excel_filepath
    )

    excel_splitter.get_docs()
    excel_splitter.get_chunks()
    excel_splitter.save_chunks_to_csv()


if __name__ == '__main__':
    main()
